import pickle
from datetime import datetime

from PIL import Image
global r
import io
import numpy as np
import os
import torch
import matplotlib.pyplot as plt
import cv2
from segment_anything.utils.transforms import ResizeLongestSide
import sys
import os
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
from service.segment.runSeg import show_mask, show_points, show_box, show_anns, generate_image_number, \
    get_resize_transform, get_bbox_from_mask, show_mask_patch
from service.segment.runSeg import Simple_coordinate_new_image, Simple_coordinate_new_image2,Simple_coordinate_new_image3
from service.segment.runSeg import load_image, load_sam_model, process_image, predict_mask,predict_mask2,prepare_image
from service.segment.runSeg import get_batched_input,batch_process_images,generate_masks
from xml.etree.ElementTree import Element, SubElement, ElementTree, tostring
from xml.dom.minidom import parseString
#持续获取前端传来的坐标数据
from routes.GlobalState import global_state
import json

def IncomingCoordinatePoint(x, y, label):
    # 从segment_route中的全局变量中获取input_point, input_label, input_stop
    input_point = global_state.input_point
    input_label = global_state.input_label
    input_stop = global_state.input_stop
    #打印当前input_point和input_label的值
    print(input_point)
    print(input_label)
    # 如果input_point或input_label为空，初始化它们为一个空列表
    if input_point is None:
        input_point = []
    if input_label is None:
        input_label = []

    # 如果input_stop为False，则可以添加点
    if not input_stop:
        input_point.append([x, y])
        input_label.append(label)
    else:
        print('此时不能添加点')

    # 将更新后的变量存储回segment_route中的全局变量中
    global_state.input_point = input_point
    global_state.input_label = input_label



# 撤回上一个点
def mouse_click_back():
    # 从segment_route中的全局变量中获取input_point, input_label, input_stop
    input_point = global_state.input_point
    input_label = global_state.input_label
    input_stop = global_state.input_stop

    # 如果input_stop为False，则可以撤销
    if not input_stop:
        input_point.pop()
        input_label.pop()
    else:
        print('此时不能撤销')

    # 将更新后的变量存储回segment_route中的全局变量中
    global_state.input_point = input_point
    global_state.input_label = input_label
    print(input_point)
    print(input_label)




# 定义模型初始化函数
def init_model(image_name):
    global sam, predictor
# 生成图片路径
    image_path = os.path.abspath(os.path.join('service', 'segment', 'image', image_name))
    print('image_path:' + image_path)
    image = load_image(image_path)
    input_point = []
    input_label = []
    # inpute_box初始化暂时为空
    input_box_arr = None
    # mask_input初始化为空
    mask_input = None
    # 是否输出多个遮罩--暂时不输出
    multimask_output = False
    input_stop = False
    sam, predictor = load_sam_model()
    process_image(predictor, image)
    return image,sam, predictor, input_point, input_label, input_box_arr, mask_input, multimask_output, input_stop,image_path


# 定义更新mask的函数
def update_mask(image, sam, predictor, input_point, input_label, mask_input, multimask_output, input_stop,
                input_box_arr):
    # 如果input_point为空，初始化它为一个空列表
    if input_point is None:
        input_point = []
    # 将input_point和input_label转换为array(如果不为空)
    if len(input_point) != 0:
        input_point_arr = np.array(input_point)
        input_label_arr = np.array(input_label)
    else:
        input_point_arr = None
        input_label_arr = None
    #print(input_box_arr)
    #input_box_arr=np.array(input_box_arr)


    # 首先进行一次预测，获取masks, scores, logits
    masks, scores, logits = predict_mask(predictor, input_point_arr, input_label_arr, multimask_output, mask_input, input_box_arr)

    # 然后更新mask_input为上一次预测的结果
    mask_input = logits[np.argmax(scores)]

    return masks, scores, logits, mask_input, input_stop,input_point_arr,input_label_arr
#保存图片函数
#生成不规则遮罩图片函数
def generate_masked_image(image, masks):
    # 将遮罩应用到图像上
    plt.figure(figsize=(10, 10))
    # 移除遮罩的不必要的维度
    masks2 = np.squeeze(masks)
    # 将masks2复制三次并将结果堆叠在一起，以匹配image的形状
    masks2 = np.stack([masks2, masks2, masks2], axis=-1)
    # 将遮罩区域的图像部分保留，而将其他部分的原图像设置为黑色
    masked_image = np.where(masks2, image, [0, 0, 0])
    # 将图像的深度转换为8位
    masked_image = cv2.convertScaleAbs(masked_image)
    # 将masked_image转换为RGBA格式
    masked_image_rgba = cv2.cvtColor(masked_image, cv2.COLOR_RGB2RGBA)
    # 将所有黑色的像素的透明度设置为0
    masked_image_rgba[(masked_image_rgba[:, :, :3] == [0, 0, 0]).all(axis=2)] = [0, 0, 0, 0]
    xmin, ymin, xmax, ymax = get_bbox_from_mask(masks)
    # 根据边界框的坐标裁剪图像
    cropped_image = masked_image_rgba[ymin:ymax, xmin:xmax]
    plt.imshow(cropped_image)
    plt.axis('off')
    # 移除图像周围的空白边缘
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    # 将图像保存为一个临时文件
    img_io = io.BytesIO()
    plt.savefig(img_io, format='png', bbox_inches='tight', pad_inches=0, transparent=True)
    img_io.seek(0)
    return img_io, xmin, ymin, xmax, ymax
#保存图片函数
def save_masked_image(img_io):
    # 在这个位置保不存在标记点的图片
    save_dir1 = os.path.abspath(os.path.join(os.path.dirname(__file__), 'SaveTheImage'))
    print("save_dir1：" + save_dir1)
    # 获取新的图片编号
    image_number = generate_image_number(save_dir1)
    if not os.path.exists(save_dir1):
        os.makedirs(save_dir1)
    # 生成保存图像的路径
    save_path1 = os.path.join(save_dir1, f'{image_number}.png')
    # 使用PIL库的Image模块来保存图像
    img = Image.open(img_io)
    img.save(save_path1, 'PNG')
    print("保存图像到指定路径：" + save_path1)
    return save_path1
#计算图片大小
def calculate_size(save_path1):
    size = os.path.getsize(save_path1)
    return size
def create_class_index_mapping():
    # 创建一个包含100个具体事物名称的列表
    class_names = [
        'dog', 'cat', 'bird', 'fish', 'elephant', 'lion', 'tiger', 'bear', 'shark', 'whale',
        'rose', 'tulip', 'daisy', 'sunflower', 'orchid', 'lily', 'cactus', 'bamboo', 'fern', 'moss',
        'man', 'woman', 'boy', 'girl', 'elderly_man', 'elderly_woman', 'baby_boy', 'baby_girl', 'teenager_boy', 'teenager_girl',
        # ...更多类别
    ]
    # 创建一个字典，其中键是类别名称，值是对应的索引
    class_name_to_index = {class_name: index for index, class_name in enumerate(class_names)}
    return class_name_to_index

#保存标注yolo格式
def save_and_add_picture_yolo(workings_toFile):
    save_dir = os.path.abspath(os.path.join('..', 'ThinkStation','service', 'segment', 'SaveTheImage', 'yolo'))
    image_number = generate_image_number(save_dir)
    #创建一个同名的文件夹用以存储图片
    image_dir = os.path.join(save_dir, f'{image_number}')
    os.makedirs(image_dir, exist_ok=True)
    # 创建YOLO格式的标注
    with open(os.path.join(image_dir, f'toYolo.txt'), 'w') as f:
        for i,workingTOFile in enumerate(workings_toFile):
            #打印workingTOFile中的所有信息
            print("workingTOFile:"+str(workingTOFile))
            xmin, ymin, xmax, ymax = workingTOFile['xmin'], workingTOFile['ymin'], workingTOFile['xmax'], workingTOFile['ymax']
            image_path = workingTOFile['pathOld']
            # 从文件中加载图像
            image = cv2.imread(image_path)
            #保存图片到指定路径
            save_image_path = os.path.join(image_dir, f'{i}.png')
            cv2.imwrite(save_image_path, image)
            boxes = [[xmin, ymin, xmax, ymax]]
            classes = [workingTOFile['classification']]
            print("classes:"+str(classes))

            for box, class_name in zip(boxes, classes):
                # 假设你已经有了一个将类别名称映射到索引的字典
                class_index = class_name
                # 计算边界框的中心点坐标和宽高（相对于图像的宽度和高度）
                x_center = (int(box[0]) + int(box[2])) / 2 / image.shape[1]
                y_center = (int(box[1]) + int(box[3])) / 2 / image.shape[0]
                width = (int(box[2]) - int(box[0])) / image.shape[1]
                height = (int(box[3]) - int(box[1])) / image.shape[0]
                # 写入文件
                f.write(f'{class_index} {x_center} {y_center} {width} {height}\n')

    # 返回txt文件存储路径和导出成功讯息
    yolo_path = os.path.join(image_dir, f'toYolo.txt')
    print("导出成功：" + yolo_path)
    return yolo_path, image_dir
#保存标注xml格式

def save_and_add_picture_xml(workings_toFile):
    save_dir = os.path.abspath(os.path.join('..', 'ThinkStation','service', 'segment', 'SaveTheImage', 'xml'))
    image_number = generate_image_number(save_dir)
    image_dir = os.path.join(save_dir, f'{image_number}')
    os.makedirs(image_dir, exist_ok=True)

    root = Element('annotations')

    for i, workingTOFile in enumerate(workings_toFile):
        xmin, ymin, xmax, ymax = workingTOFile['xmin'], workingTOFile['ymin'], workingTOFile['xmax'], workingTOFile[
            'ymax']
        image_path = workingTOFile['pathOld']
        image = cv2.imread(image_path)
        save_image_path = os.path.join(image_dir, f'{i}.png')
        cv2.imwrite(save_image_path, image)
        boxes = [[xmin, ymin, xmax, ymax]]
        classes = [workingTOFile['classification']]

        for box, class_name in zip(boxes, classes):
            box_element = SubElement(root, 'box')
            class_index = class_name
            class_index_element = SubElement(box_element, 'class_index')
            class_index_element.text = str(class_index)

            x_center = (int(box[0]) + int(box[2])) / 2 / image.shape[1]
            y_center = (int(box[1]) + int(box[3])) / 2 / image.shape[0]
            width = (int(box[2]) - int(box[0])) / image.shape[1]
            height = (int(box[3]) - int(box[1])) / image.shape[0]

            x_center_element = SubElement(box_element, 'x_center')
            x_center_element.text = str(x_center)

            y_center_element = SubElement(box_element, 'y_center')
            y_center_element.text = str(y_center)

            width_element = SubElement(box_element, 'width')
            width_element.text = str(width)

            height_element = SubElement(box_element, 'height')
            height_element.text = str(height)

    tree = ElementTree(root)
    xml_path = os.path.join(image_dir, f'{image_number}.xml')
    tree.write(xml_path)
    print("导出成功："+xml_path)
    return xml_path,image_dir
#保存标注coco格式


def save_and_add_picture_coco(workings_toFile):
    save_dir = os.path.abspath(os.path.join('..', 'ThinkStation','service', 'segment', 'SaveTheImage', 'coco'))
    image_number = generate_image_number(save_dir)
    # 创建一个同名的文件夹用以存储图片
    image_dir = os.path.join(save_dir, f'{image_number}')
    os.makedirs(image_dir, exist_ok=True)

    # 创建COCO格式的标注
    coco_annotations = {
        "images": [],
        "annotations": [],
        "categories": []
    }

    for i, workingTOFile in enumerate(workings_toFile):
        xmin, ymin, xmax, ymax = workingTOFile['xmin'], workingTOFile['ymin'], workingTOFile['xmax'], workingTOFile[
            'ymax']


        classification = workingTOFile['classification']
        boxes = [[xmin, ymin, xmax, ymax]]
        classes = [classification]  # 目标的类别

        for box, class_id in zip(boxes, classes):
            # 从文件中加载图像
            image_path = workingTOFile['pathOld']
            image = cv2.imread(image_path)
            # 保存图片到指定路径
            save_image_path = os.path.join(image_dir, f'{i}.png')
            cv2.imwrite(save_image_path, image)

            # 添加图像信息
            coco_annotations["images"].append({
                "id": i,
                "width": image.shape[1],
                "height": image.shape[0],
                "file_name": f"{i}.png"
            })
            # 计算边界框的中心点坐标和宽高（相对于图像的宽度和高度）
            x_center = (int(box[0]) + int(box[2])) / 2 / image.shape[1]
            y_center = (int(box[1]) + int(box[3])) / 2 / image.shape[0]
            width = (int(box[2]) - int(box[0])) / image.shape[1]
            height = (int(box[3]) - int(box[1])) / image.shape[0]

            # 添加标注信息
            coco_annotations["annotations"].append({
                "id": i,
                "image_id": i,
                "category_id": class_id,
                "bbox": [x_center, y_center, width, height],
                "area": width * height,
                "iscrowd": 0
            })

    # 添加类别信息
    for class_name, class_index in create_class_index_mapping().items():
        coco_annotations["categories"].append({
            "id": class_index,
            "name": class_name
        })

    # 将标注信息写入到文件中
    coco_path = os.path.join(image_dir, f'{image_number}.json')
    with open(coco_path, 'w') as f:
        json.dump(coco_annotations, f)
    print("导出成功：" + coco_path)

    return coco_path,image_dir


# 保存标注voc    格式
def save_and_add_picture_voc(workings_toFile):
    save_dir = os.path.abspath(os.path.join('..', 'ThinkStation', 'service', 'segment', 'SaveTheImage', 'voc'))
    image_number = generate_image_number(save_dir)
    # 创建一个同名的文件夹用以存储图片
    image_dir = os.path.join(save_dir, f'{image_number}')
    os.makedirs(image_dir, exist_ok=True)

    for i, workingTOFile in enumerate(workings_toFile):
        xmin, ymin, xmax, ymax = workingTOFile['xmin'], workingTOFile['ymin'], workingTOFile['xmax'], workingTOFile[
            'ymax']
        image_path = workingTOFile['pathOld']
        class_name = workingTOFile['classification']

        image = cv2.imread(image_path)
        save_image_path = os.path.join(image_dir, f'{i}.png')  # 修改这里，将图片保存到新创建的文件夹中
        cv2.imwrite(save_image_path, image)

        annotation = Element('annotation')
        SubElement(annotation, 'folder').text = 'VOC2007'
        SubElement(annotation, 'filename').text = f'{i}.png'
        SubElement(annotation, 'path').text = save_image_path
        source = SubElement(annotation, 'source')
        SubElement(source, 'database').text = 'Unknown'
        size = SubElement(annotation, 'size')
        SubElement(size, 'width').text = str(image.shape[1])
        SubElement(size, 'height').text = str(image.shape[0])
        SubElement(size, 'depth').text = str(image.shape[2])
        SubElement(annotation, 'segmented').text = '0'

        obj = SubElement(annotation, 'object')
        SubElement(obj, 'name').text = class_name
        SubElement(obj, 'pose').text = 'Unspecified'
        SubElement(obj, 'truncated').text = '0'
        SubElement(obj, 'difficult').text = '0'
        bndbox = SubElement(obj, 'bndbox')
        SubElement(bndbox, 'xmin').text = str(xmin)
        SubElement(bndbox, 'ymin').text = str(ymin)
        SubElement(bndbox, 'xmax').text = str(xmax)
        SubElement(bndbox, 'ymax').text = str(ymax)

        tree = ElementTree(annotation)
        voc_path = os.path.join(image_dir, f'{i}.xml')
        tree.write(voc_path)
        print("导出成功：" + voc_path)

    return voc_path, image_dir
